#pragma once

#include <ros/ros.h>
#include <std_msgs/UInt8.h>
#include <std_msgs/Int8.h>
#include <unitree_legged_msgs/LowCmd.h>
#include <unitree_legged_msgs/LowState.h>
#include "stair_ctrl/JoyStick.h"
#include "stair_ctrl/NN.h"
#include "unitree_legged_sdk/unitree_legged_sdk.h"
#include "unitree_legged_sdk/joystick.h"

#include "convert.h"
#include "joy_stick_advanced.hpp"
#include "state_machine.hpp"
#include "interface.hpp"
#include "step_plan_interface.hpp"
#include "NN_controller.hpp"

#include <iostream>
#include <Eigen/Core>
#include <ctime>

using namespace UNITREE_LEGGED_SDK;

class stairCtrl
{
private:
    unitree_legged_msgs::LowCmd low_cmd_ros;
    unitree_legged_msgs::LowState low_state_ros;
    stair_ctrl::JoyStick joy_stick;
    stair_ctrl::NN nn_msg;

    std::string low_state_topic;
    ros::Publisher low_cmd_pub;
    ros::Publisher joy_stick_state_pub;
    ros::Publisher reset_pub;
    ros::Publisher state_machine_pub;
    ros::Publisher nn_pub;
    ros::Subscriber low_state_sub;

    ros::WallTimer ctrl_timer;

    int rate = 100;
    int powerLevel = 1;
    double kp_final, kd_final, kp_init, kd_init, kd_low, kd_high;
    xRockerBtnDataStruct key_data;
    joyStickAdvanced joy_stick_advanced;

    Eigen::VectorXd joint_nominal_position;
    std::array<int, 12> joint_idx = {FR_0, FR_1, FR_2,
                                    FL_0, FL_1, FL_2,
                                    RR_0, RR_1, RR_2,
                                    RL_0, RL_1, RL_2};

    stateMachine state_machine;
    bool state_changed = false;
    double stand_seconds = 3.0;
    double stand_count = 0.0;

    robotInterface<float> robot_interface;
    cmdInterface<float> cmd_interface;
    stepPlanInterface<float> step_plan_interface;
    NNController nn_ctrl;
    std::string nn_log_path;
    
public:
    stairCtrl(ros::NodeHandle *nh,
              std::string low_state_topic = "/ros_udp_comm_low/low_state",
              std::string low_cmd_topic = "low_cmd",
              std::string joy_stick_topic = "joy_stick",
              std::string reset_topic = "reset",
              std::string state_machine_topic = "state",
              std::string nn_state_topic = "nn_state",
              std::string nn_param_path = "",
              std::string nn_log_prefix = "")
    {
        auto hints = ros::TransportHints().tcpNoDelay();
        
        low_cmd_pub = nh->advertise<unitree_legged_msgs::LowCmd>(low_cmd_topic, 10);
        joy_stick_state_pub = nh->advertise<stair_ctrl::JoyStick>(joy_stick_topic, 10);
        reset_pub = nh->advertise<std_msgs::UInt8>(reset_topic, 10);
        state_machine_pub = nh->advertise<std_msgs::Int8>(state_machine_topic, 10);
        nn_pub = nh->advertise<stair_ctrl::NN>(nn_state_topic, 10);
        
        low_state_sub = nh->subscribe(low_state_topic, 10, &stairCtrl::lowStateCallback, this, hints);

        ctrl_timer = nh->createWallTimer(ros::WallDuration(1.0 / rate), &stairCtrl::periodicControlCallback, this);
        ctrl_timer.stop();

        init();
        initNN(nn_param_path);

        nn_log_path = nn_log_prefix;
        if (nn_log_path != "")
        {
            std::time_t now_ = time(0);
            std::tm *ltm = localtime(&now_);
            std::stringstream logfile;
            logfile << nn_log_path << "/"
                    << 1900 + ltm->tm_year << "_" 
                    << 1 + ltm->tm_mon << "_" 
                    << ltm->tm_mday << "_" 
                    << ltm->tm_hour << "_" 
                    << ltm->tm_min << "_" 
                    << ltm->tm_sec << ".csv";
            nn_log_path = logfile.str();
        }

        step_plan_interface.init(nh);
    }

    void init()
    {
        joint_nominal_position.setZero(12);
        joint_nominal_position << 0.0, 0.79, -1.58,
            0.0, 0.79, -1.58,
            0.0, 0.79, -1.58,
            0.0, 0.79, -1.58;
        kp_final = 40.0;
        kd_final = 1.0;
        kp_init = 5.0;
        kd_init = 1.0;
        kd_low = 1.0;
        kd_high = 3.0;

        joy_stick_advanced = joyStickAdvanced();
        state_machine = stateMachine();
        robot_interface = robotInterface<float>();
        cmd_interface = cmdInterface<float>();
        step_plan_interface = stepPlanInterface<float>(rate, 0.5);

        robot_interface.setPD(static_cast<float>(kp_final), static_cast<float>(kd_final));
        cmd_interface.setCmdRange(1.0, 0.4, 0.8);

        low_cmd_ros.head[0] = 0xFE;
        low_cmd_ros.head[1] = 0xEF;
        low_cmd_ros.levelFlag = LOWLEVEL;
    }

    void initNN(std::string nn_param_path)
    {
        nn_ctrl = NNController(nn_param_path);
    }

    void lowStateCallback(const unitree_legged_msgs::LowState::ConstPtr &msg)
    {
        low_state_ros = *msg.get();
        memcpy(&key_data, &msg->wirelessRemote[0], 40);
        joy_stick_advanced.getJoyStickRosMsg(key_data, joy_stick);
        joy_stick_state_pub.publish(joy_stick);
    }

    void publishState()
    {
        std_msgs::Int8 msg = std_msgs::Int8();
        msg.data = state_machine.getState();
        state_machine_pub.publish(msg);
    }
    
    void control()
    {
        switch (state_machine.getState())
        {
            case -2:
                damping(kd_low);
                break;
            case -1:
                damping(kd_high);
                break;
            case 0:
                normalState();
                break;
            case 1:
                standUp();
                break;
            case 2:
                standing();
                break;
            case 3:
                nnCtrl();
                break;

            default:
                break;
        }
    }

    void damping(double kd)
    {
        for (int i = 0; i < 12; ++i)
        {
            low_cmd_ros.motorCmd[joint_idx[i]].mode = 0x0A;
            low_cmd_ros.motorCmd[joint_idx[i]].q = PosStopF;
            low_cmd_ros.motorCmd[joint_idx[i]].dq = 0.0f;
            low_cmd_ros.motorCmd[joint_idx[i]].Kp = 0.0f;
            low_cmd_ros.motorCmd[joint_idx[i]].Kd = static_cast<float>(kd);
            low_cmd_ros.motorCmd[joint_idx[i]].tau = 0.0f;
        }
        low_cmd_pub.publish(low_cmd_ros);
        stand_count = 0.0;
    }

    void normalState()
    {
        if (joy_stick_advanced.joy_stick.keys.B == 3)
        {
            std_msgs::UInt8 msg = std_msgs::UInt8();
            msg.data = 233;
            reset_pub.publish(msg);
        }
        stand_count = 0.0;
    }

    void standUp()
    {
        stand_count += 1.0;
        double ratio = stand_count / (stand_seconds * rate);
        float kp = static_cast<float>(kp_init * (1 - ratio) + kp_final * ratio);
        float kd = static_cast<float>(kd_init * (1 - ratio) + kd_final * ratio);
        for (int i = 0; i < 12; ++i)
        {
            low_cmd_ros.motorCmd[joint_idx[i]].mode = 0x0A;
            low_cmd_ros.motorCmd[joint_idx[i]].q = joint_nominal_position[i];
            low_cmd_ros.motorCmd[joint_idx[i]].dq = 0.0f;
            low_cmd_ros.motorCmd[joint_idx[i]].Kp = kp;
            low_cmd_ros.motorCmd[joint_idx[i]].Kd = kd;
            low_cmd_ros.motorCmd[joint_idx[i]].tau = 0.0f;
        }
        low_cmd_pub.publish(low_cmd_ros);
        // std::cout << "Current mode: [StandUp]; count: " << stand_count << "/" << stand_seconds * rate << std::endl;

        if (stand_count > stand_seconds * rate)
        {
            state_machine.setState(2);
            std::cout << "Switch to state [2]!!! Press Start for NN ctrl!!!" << std::endl;
        }
    }

    void standing()
    {
        for (int i = 0; i < 12; ++i)
        {
            low_cmd_ros.motorCmd[joint_idx[i]].q = joint_nominal_position[i];
            low_cmd_ros.motorCmd[joint_idx[i]].dq = 0.0f;
            low_cmd_ros.motorCmd[joint_idx[i]].Kp = kp_final;
            low_cmd_ros.motorCmd[joint_idx[i]].Kd = kd_final;
            low_cmd_ros.motorCmd[joint_idx[i]].tau = 0.0f;
        }
        low_cmd_pub.publish(low_cmd_ros);

        step_plan_interface.blind = true;
    }

    void nnCtrl()
    {
        updateRobotInterface();
        updateCmdInterface();
        updateStepPlanInterface();
        
        nn_ctrl.update(robot_interface, cmd_interface, step_plan_interface);
        robot_interface.simpleWcProtect(23.7f, 30.1f, 4.5f);
        
        setRosLowCmd();
        low_cmd_pub.publish(low_cmd_ros);

        nn_msg = stair_ctrl::NN();
        auto input = nn_ctrl.getNNInput();
        auto output = nn_ctrl.getNNOutput();
        auto input_raw = nn_ctrl.getNNInputRaw();
        for (int i = 0; i < 62; ++i)
        {
            nn_msg.input[i] = input[i];
        }
        for (int i = 0; i < 62; ++i)
        {
            nn_msg.input_raw[i] = input_raw[i];
        }
        for (int i = 0; i < 12; ++i)
        {
            nn_msg.output[i] = output[i];
        }
        nn_pub.publish(nn_msg);

        if (nn_log_path != "")
        {
            nn_ctrl.updateLog(nn_log_path);
        }
    }

    void updateRobotInterface()
    {
        for (int i = 0; i < 12; ++i)
        {
            robot_interface.joint_pos[i] = low_state_ros.motorState[joint_idx[i]].q;
            robot_interface.joint_vel[i] = low_state_ros.motorState[joint_idx[i]].dq;
        }
        robot_interface.joint_vel_filtered = robot_interface.joint_vel_filtered * (1.0 - robot_interface.joint_vel_filter_alpha) + robot_interface.joint_vel * robot_interface.joint_vel_filter_alpha;

        robot_interface.body_quat[0] = low_state_ros.imu.quaternion[0];
        robot_interface.body_quat[1] = low_state_ros.imu.quaternion[1];
        robot_interface.body_quat[2] = low_state_ros.imu.quaternion[2];
        robot_interface.body_quat[3] = low_state_ros.imu.quaternion[3];

        robot_interface.body_ang_vel[0] = low_state_ros.imu.gyroscope[0];
        robot_interface.body_ang_vel[1] = low_state_ros.imu.gyroscope[1];
        robot_interface.body_ang_vel[2] = low_state_ros.imu.gyroscope[2];
        
        robot_interface.body_ang_pos[0] = low_state_ros.imu.rpy[0];
        robot_interface.body_ang_pos[1] = low_state_ros.imu.rpy[1];
        robot_interface.body_ang_pos[2] = low_state_ros.imu.rpy[2];

        robot_interface.foot_force[0] = low_state_ros.footForce[0];
        robot_interface.foot_force[1] = low_state_ros.footForce[1];
        robot_interface.foot_force[2] = low_state_ros.footForce[2];
        robot_interface.foot_force[3] = low_state_ros.footForce[3];

        robot_interface.body_lin_acc[0] = low_state_ros.imu.accelerometer[0];
        robot_interface.body_lin_acc[1] = low_state_ros.imu.accelerometer[1];
        robot_interface.body_lin_acc[2] = low_state_ros.imu.accelerometer[2];
    }

    void setRosLowCmd()
    {
        for (int i = 0; i < 12; ++i)
        {
            low_cmd_ros.motorCmd[joint_idx[i]].q = robot_interface.joint_pos_des[i];
            low_cmd_ros.motorCmd[joint_idx[i]].dq = 0.0f;
            low_cmd_ros.motorCmd[joint_idx[i]].Kp = robot_interface.joint_kp[i];
            low_cmd_ros.motorCmd[joint_idx[i]].Kd = robot_interface.joint_kd[i];
            low_cmd_ros.motorCmd[joint_idx[i]].tau = 0.0f;
        }
    }

    void updateCmdInterface()
    {
        cmd_interface.cmd[0] = joy_stick_advanced.joy_stick.ly * cmd_interface.cmd_range[0];
        cmd_interface.cmd[1] = -joy_stick_advanced.joy_stick.lx * cmd_interface.cmd_range[1];
        cmd_interface.cmd[2] = -joy_stick_advanced.joy_stick.rx * cmd_interface.cmd_range[2];
    }

    void updateStepPlanInterface()
    {
        step_plan_interface.update(robot_interface, cmd_interface);
        if (joy_stick_advanced.joy_stick.keys.L1 == 3)
        {
            step_plan_interface.switchMode();
        }
    }

    void periodicControlCallback(const ros::WallTimerEvent& event)
    {
        joy_stick_advanced.update(key_data);
        state_changed = state_machine.updateState(joy_stick_advanced);

        publishState();
        control();
    }

    void start()
    {
        // initialize
        ros::Rate loop_rate(rate);
        for (int i = 0; i < rate * 5; ++i)
        {
            loop_rate.sleep();
            ros::spinOnce(); // this spin may cause the joystick init problem

            joy_stick_advanced.update(key_data);
        }

        ctrl_timer.start();
        
        ros::spin();
    }

};